from 搜索二叉树 import BitTree, BST

class AVLTree(BitTree):
    def __init__(self, val):
        super().__init__(val)
        self.bf = 0

    

class AVLTree(BST):
    def __init__(self, li=None):
        super().__init__(li)
    
    def rotate_left(self, p, c):
        s2 = c.lchild 
        p.rchild = s2
        if s2:
            s2.parent = p
        p.parent = c 
        c.lchild = p
        p.bf = 0
        c.bf = 0
    
    def rotate_right(self, p, c):
        s2 = c.rchild
        p.lchild = s2
        if s2:
            s2.parent = p
        p.parent = c
        c.rchild = p
        p.bf = 0
        c.bf = 0
    
    def rotate_right_left(self, p, c):
        g = c.lchild
        s3 = g.rchild
        c.lchild = s3
        if s3:
            s3.parent = c
        g.rchild = c
        c.parent = g
        
        s2 = g.lchild
        p.rchild = s2
        if s2:
            s2.parent = p
        g.lchild = p
        p.parent =g
        if g.bf > 0:
            p.bf = -1
            c.bf = 0
        elif g.bf < 0:
            p.bf = 0
            c.bf = 1
        else:
            p.bf = 0
            c.bf = 0
    
    def rotate_left_right(self, p, c):
        g = c.rchild
        s2 = g.lchild
        c.rchild = s2
        if s2:
            s2.parent = c
        g.lchild = c
        c.parent = g

        s3 = g.rchild
        p.lchild = s3
        if s3:
            s3.parent = p
        p.parent = g
        g.rchild = p

        if g.bf > 0:
            c.bf = -1
            p.bf = 0
        elif g.bf < 0:
            c.bf = 0
            p.bf = 1
        else:
            c.bf = 0
            p.bf = 0

        



            


